function [optima, RegionSet, SampleSet, RegionSampleId] = PRS_MMO_v4_mergesort( Problem, AlgorithmM )
% PRS_MMO: Partition-based Random Search for Multimodal Optimization for 
%   minimization problems.
% decision variable: d dimensions; feature of regions: m dimensions.
%INPUT
%   Problem.
%       Dimension: scalar, the dimension of the decision variables
%       Domain: 1 -by- m vector, features of the domain
%       Partition: function handle of partitioning strategy
%           [RegionNum, Region, RegionSampleId] = Partition(Region, SampleSet)
%               RegionNum: scalar, number of new partitioned regions
%               Region: RegionNum -by- (2+m) matrix,
%                   [partition depth, partitionable, features ...]
%                   partition depth = 0 for the original domain
%                   partitionable: scalar, 1: partitionable, 0: non-partitionable
%               RegionSampleId: RegionNum -by- 1 cell, the sample id that
%                   belongs to each region
%               SampleSet: [SampleID, objective value, x1, x2, ...]
%       Sampling: function handle of sampling in a particular region
%           [Sample] = Sampling(n, region)
%               n: number of samples
%               region: 1 -by- m vector, [features...]
%               Sample: n -by- (1+d) matrix, [objective value, x1, x2, ...]
%   AlgorithmM.
%       n0: scalar, the base sample size
%       QuantileLevel: scalar, <0.5, the quantile level of promisinhg region definition
%       NewBudget: scalar, added budget size at each stage
%       MaxSampleSize: scalar, maximal sample size in one region
%       Radius: scalar, the radius distance that be regarded as neighbor
%       StopCriteria: [X,Y] (the maximal budget size is 1e6)
%           [1,Y]: reach maximum budget Y
%           [2,Y]: reach the required number of optima Y
%           [3,Y]: the number of optima is not changed in Y iterations
%   OptSet.
%       sol: ~ -by- d matrix, the optimal solution
%       fval: scalar, the global optimum
%       epsilon: scalar, the difference of the value considered to be optimum
%       radius: scalar, the distance of the solution considered to be found
%
%OUTPUT
%   optima: ~ -by- (2+d) matrix, the final optimal solution set,
%         [sample id, objective value, x1, x2, ...]
%   RegionSet: ~ -by- (2+m) matrix, the whole subregion set
%         [partition depth, partitionable, features ...]
%   SampleSet: ~ -by- (1+d) matrix, the original sample set
%         [objective value, x1, x2, ...]
%   RegionSampleId: ~ -by- 1 cell, the id of the samples in each region

%% INITIALIZATION
% problem setting
d = Problem.Dimension;
domain = [0, 1, Problem.Domain]; % [partition depth, partitionable, features ...]
Partition = Problem.Partition;
Sampling = Problem.Sampling;
% algorithm setting
n0 = AlgorithmM.n0;
alpha = AlgorithmM.QuantileLevel;
Delta = AlgorithmM.NewBudget;
MaxSampleSize = AlgorithmM.MaxSampleSize;
StopCriteria = AlgorithmM.StopCriteria;
N = (StopCriteria(1)==1)*StopCriteria(2) + (StopCriteria(1)~=1)*1e6; % the maximal budget size
radius = AlgorithmM.radius;
% algorithm archive
RegionSet = zeros(round(N/n0), length(domain));  % [partition depth, partitionable, features ...]
RegionSet(1,:) = domain;
RegionsNum = 1;  % the number of exist regions
PartitionableNum = 1;  % the number of exist partitionable regions
CurrentMaxPartitionDepth = 0;  % the maximum depth of partitioning at current iteration
RegionSampleId = cell(round(N/n0), 1);  % id of group samples
SampleSet = zeros(N, 2+d);  % [SampleID, objective value, x1, x2, ...]
SampleSet(:,1) = 1:size(SampleSet,1);
SampleRank = zeros(N,1);  % the rank of each sample, 1: the worst one
ranking = [];  %  the solution id from the worst to the best
N0 = 0;  % budget allocated
N0_checked = 0;  % number of samples that have been checked by the optima set
OptCandidate = [zeros(N,1), ones(N,1)];
% [1: the potential optima, 1: not be dominated by neibours]
GroupInf = zeros(round(N/n0), 5); % [size, mean, std, modified size, weight]
%% OPTIMIZATION
NumOpt = 0;  % the current optimal value
NumIterationRepeat = 0;  % the number of iterations that the number of optima is not changed
PartitionList = zeros(Delta,1);  % the list of regions to be partitioined
PartitionList(1) = 1;
PartitionListNum = 1;  % the number of the partition list
WeightChangedList = zeros(round(N/n0),1);  % the list of regions whose weight is changed
WeightChangedListNum = 0;  % the number of the region list whose weight is changed
all_weight_changed = 0;  % if all weights should be recalculated
NewOptCandidate = 0;  % if there are new regions reach the non-partitionable size
BestGroupInf = [1,inf,0];  % the best group information
%       [group id, quantile, modified sample size]
optima = [];  % the set of optima
% subsequent iterations
while N0 < N && PartitionableNum>0 ...
    && (StopCriteria(1)~=2 || NumOpt<StopCriteria(2)) ...
    && (StopCriteria(1)~=3 || NumOpt==0 || NumIterationRepeat<StopCriteria(2))
%% partition partitionable regions
    i = 1;
    while i <= PartitionListNum
        % partition
        iPartitionRegion = PartitionList(i);
        [NewRegionNum, NewRegion, NewRegionSampleId]...
            = Partition(RegionSet(iPartitionRegion,:), SampleSet(RegionSampleId{iPartitionRegion},:));
        NewRegionId = [iPartitionRegion,(RegionsNum+1):(RegionsNum+NewRegionNum-1)]; % id of new-partitioned regions
        RegionSet(NewRegionId,:) = NewRegion;
        RegionSampleId(NewRegionId) = NewRegionSampleId;
        RegionsNum = RegionsNum+NewRegionNum-1;  % update total region number
        PartitionableNum = PartitionableNum-1+sum(NewRegion(:,2));  % update partitionable region number
        % update the maximal partition depth
        MaxPartitionDepth_temp = max(NewRegion(:,1));
        if MaxPartitionDepth_temp>CurrentMaxPartitionDepth
            CurrentMaxPartitionDepth = MaxPartitionDepth_temp;
            GroupInf(1:RegionsNum,4) = max(2,round(GroupInf(1:RegionsNum,1).*RegionSet(1:RegionsNum,1)./CurrentMaxPartitionDepth));
            all_weight_changed = 1;
            BestGroupInf(3) = GroupInf(BestGroupInf(1),4);
        end
        % first stage sampling and group information updating
        for j = NewRegionId
            if RegionSet(j,2)==0
                %% non-partitionable region
                if isempty(RegionSampleId{j})
                    N0 = N0+1;
                    SampleSet(N0,2:end) = Sampling(1,RegionSet(j,3:end));
                    RegionSampleId{j} = N0;
                end
                GroupInf(j,1) = length(RegionSampleId{j});
                group_update(j)
                % all samples in this region are potential optima
                OptCandidate(RegionSampleId{j})=1;  % update the candidate optimal solution
                NewOptCandidate = 1;
            else
                %% partitionable region
                GroupInf(j,1) = length(RegionSampleId{j});  % update the sample size
                if GroupInf(j,1)>=MaxSampleSize
                    % if the sample size in the new region also reaches the threshold
                    PartitionListNum = PartitionListNum+1;
                    PartitionList(PartitionListNum) = j;
                else
                    if GroupInf(j,1)<n0
                    % if the sample size in the new region is not enough
                        Add_n = min([n0-GroupInf(j,1),N-N0]);
                        NewSampleId = ((N0+1):(N0+Add_n))';
                        SampleSet(NewSampleId,2:end) = Sampling(Add_n,RegionSet(j,3:end));
                        N0 = N0+Add_n;
                        RegionSampleId{j} = [RegionSampleId{j};NewSampleId];
                        GroupInf(j,1) = GroupInf(j,1)+Add_n;
                    end
                    group_update(j)
                    WeightChangedListNum = WeightChangedListNum+1;
                    WeightChangedList(WeightChangedListNum) = j;
                end
            end
        end
        i = i+1;
    end
    PartitionListNum = 0;
%% update the optima set
    if StopCriteria(1)~=1 && NewOptCandidate == 1
        optima = extract();
        NumOpt_old = NumOpt;
        NumOpt = size(optima,1);
        if NumOpt_old ~= NumOpt
            NumIterationRepeat = 0;
        end
        NewOptCandidate = 0;
    end
    NumIterationRepeat = NumIterationRepeat+1;
%% update the weights using baqm
    if all_weight_changed == 1
    % if the maximal partition depth or the minimal quantile is changed
        WeightChangedListNum = PartitionableNum;
        WeightChangedList(1:PartitionableNum) = find(RegionSet(1:RegionsNum,2)==1)';
        all_weight_changed = 0;
    end
    baqm(WeightChangedList(1:WeightChangedListNum))
    WeightChangedListNum = 0;
%% additional budget
    Add_n = min([N-N0,Delta]);  % new budget size
    [ n2 ] = add_budget(Add_n);  % K -by- 2 matrix, [region id, new sample size]
    for i = 1:size(n2,1)
        iregion = n2(i,1);  % region id
        NewSampleId = ((N0+1):(N0+n2(i,2)))';  % sample id
        SampleSet(NewSampleId,2:end) = Sampling(n2(i,2),RegionSet(iregion,3:end));
        N0 = N0+n2(i,2);
        RegionSampleId{iregion} = [RegionSampleId{iregion};NewSampleId];
        GroupInf(iregion,1) = GroupInf(iregion,1)+n2(i,2);
        if GroupInf(iregion,1)>=MaxSampleSize && RegionSet(iregion,2)==1
            PartitionListNum = PartitionListNum+1;
            PartitionList(PartitionListNum) = iregion;
        else
            group_update(iregion)
            WeightChangedListNum = WeightChangedListNum+1;
            WeightChangedList(WeightChangedListNum) = iregion;
        end
    end
end
%% final results
if StopCriteria(1)==1
    [optima] = extract();
end
RegionSet((RegionsNum+1):end, :) = [];
SampleSet((N0+1):end, :) = [];
SampleSet(:,1)=[];
RegionSampleId((RegionsNum+1):end, :) = [];

%% Plots for 2 dimension box constraint problems
% if d==2
%     figure()
%     hold on
%     scatter(SampleSet(:,3),SampleSet(:,4),2,SampleSet(:,2),'filled');
%     for i = 1:RegionsNum
%         rectangle('Position',[RegionSet(i,[3,5]),RegionSet(i,[4,6])-RegionSet(i,[3,5])])
%     end
%     for i = 1:size(optima,1)
%         plot(optima(:,3),optima(:,4),'x')
%     end
%     hold off
% end


%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% FUNCTIONS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function baqm(list)
    %baqm: calculate the weight for regions in the list using Budget
    %Allocation for Quantile Minimization
        if ~isempty(list)
            Wk_best = (1+norminv(alpha,0,1)^2-1/BestGroupInf(3));
            K = length(list); % number of groups
            cv_hat_inv = (GroupInf(list,2)-BestGroupInf(2))./GroupInf(list,3);
            Wk = (1+(cv_hat_inv.^2)-1./GroupInf(list,4));
            Ckb = Wk_best./Wk;
            P_kb = fcdf(Ckb, GroupInf(list,4)-1, ones(K,1).*BestGroupInf(3)-1); % the probability that cv_k>cv_b
            weights = P_kb./(1-P_kb); % Ni / Nb
            GroupInf(list,5) = weights;
        end
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function [ n2 ] = add_budget(NewBudget)
    % add_budget: allocate NewBudget according to the weights
    % n2: K -by- 2 matrix, [region id, new sample size]
        partitionable_list = find(RegionSet(1:RegionsNum,2));
        n1 = GroupInf(partitionable_list,4);
        N_tot = sum(n1)+NewBudget;
        n = N_tot.*GroupInf(partitionable_list,5)./sum(GroupInf(partitionable_list,5));
        n2 = zeros(PartitionableNum,1);
        while NewBudget>0
            [n_delta, k] = max(n-n1-n2);
            if n_delta<1
                add = 1;
            else
                add = min(NewBudget,round(n_delta));
            end
            n2(k) = n2(k)+add;
            NewBudget = NewBudget-add;
        end
        n2_id = find(n2>0);
        n2 = [partitionable_list(n2_id), n2(n2_id)];
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function group_update( groupId )
    % update the group information
        Samplevalue_temp = SampleSet(RegionSampleId{groupId},2);
        % update group information
        GroupInf(groupId,2) = mean(Samplevalue_temp);  % group mean
        GroupInf(groupId,3) = sqrt(var(Samplevalue_temp));  % group variance
        GroupInf(groupId,4) = max([2,round(GroupInf(groupId,1)*RegionSet(groupId,1)/CurrentMaxPartitionDepth)]);  % modified sample size
        GroupInf(groupId,5) = 0;  % the weight
        new_quantile = GroupInf(groupId,2)+GroupInf(groupId,3)*norminv(alpha,0,1);
        if groupId==BestGroupInf(1) || new_quantile<BestGroupInf(2)
            if new_quantile<=BestGroupInf(2)
                BestGroupInf = [groupId, new_quantile, GroupInf(groupId,4)];
            else
                [BestGroupInf(2),BestGroupInf(1)] = min(GroupInf(1:RegionsNum,2)+GroupInf(1:RegionsNum,3).*norminv(alpha,0,1));
                BestGroupInf(3) = GroupInf(BestGroupInf(1),4);
            end
            all_weight_changed = 1;
        end
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function [ Optima ] = extract()
    %Nitching: obtain the final optima set
        %% ranking the solutions
        merge_sort(N0_checked)
        SampleRank(ranking) = (1:N0)';
        OptimaSet = OptCandidate(ranking,1);
        % solutions have not be checked by existed optima
        NewSampleId_f = ((N0_checked+1):N0)';  % new samples that have not be compared with old optima
        NewSampleSize_f = N0-N0_checked;
        %% clearing
        sampleid_temp = find(OptimaSet,1);
        sampleid = ranking(sampleid_temp);
        while 1
            if ~isempty(optima) && ismember(SampleSet(sampleid,1),optima(:,1))
                Distance = (sqrt(sum((repmat(SampleSet(sampleid,3:end),[NewSampleSize_f,1])-SampleSet(NewSampleId_f,3:end)).^2,2)))';
                neighbor = [NewSampleId_f(Distance<=radius);sampleid];
            else
                Distance = (sqrt(sum((repmat(SampleSet(sampleid,3:end),[N0,1])-SampleSet(1:N0,3:end)).^2,2)))';
                neighbor = find(Distance<=radius);
            end
            [~,best_neighbor] = min(SampleSet(neighbor,2));
            best_neighbor = neighbor(best_neighbor);
            OptCandidate(neighbor,1) = 0; % delete all neighbor
            OptCandidate(best_neighbor,1) = OptCandidate(best_neighbor,2);  % save the best neighbor to the optima set
            OptCandidate(neighbor,2) = OptCandidate(neighbor,1);
            OptimaSet(SampleRank(neighbor)) = 0; % delete all neighbor
            OptimaSet(SampleRank(best_neighbor)) = OptCandidate(best_neighbor,1);
            next = find(OptimaSet((sampleid_temp+1):end),1);
            if isempty(next)
                break
            else
                sampleid_temp = sampleid_temp + next;
                sampleid = ranking(sampleid_temp);
            end
        end
        Optima = SampleSet(OptCandidate(1:N0,1)==1,:);
        N0_checked = N0;
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function merge_sort(N_sorted)
        if N_sorted > 0
            rank_temp1 = ranking;
            ranking = zeros(N0,1);
            [~,rank_temp2] = sort(SampleSet((N_sorted+1):N0,2),'descend');
            rank_temp2 = rank_temp2+N_sorted;
            N_nosorted = N0-N_sorted;
            ii = 1;
            jj = 1;
            kk = 1;
            while ii<=N_sorted && jj<=N_nosorted
                if SampleSet(rank_temp1(ii),2)>SampleSet(rank_temp2(jj),2)
                    ranking(kk) = rank_temp1(ii);
                    ii = ii+1;
                    kk = kk+1;
                else
                    ranking(kk) = rank_temp2(jj);
                    jj = jj+1;
                    kk = kk+1;
                end
            end
            if ii>N_sorted
                ranking(kk:end) = rank_temp2(jj:N_nosorted);
            else
                ranking(kk:end) = rank_temp1(ii:N_sorted);
            end
        else
            [~,ranking] = sort(SampleSet(1:N0,2),'descend');
        end
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
end

